#!/bin/bash
#
# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

. ../lib/shflags/shflags

SCRIPT_BASE="$(dirname "$0")/../pack_dist"
. "${SCRIPT_BASE}/common.sh"

#----------------------------------------------------------------------
# Variables to control mock behavior.

# MOCK_WPSW_*: 1 for write-protected, 0 for no, empty for error.
MOCK_WPSW_CUR=
MOCK_WPSW_BOOT=

# MOCK_*_WP: Empty = "not enabled", otherwise "enabled"
MOCK_MAIN_WP=
MOCK_EC_WP=
MOCK_PD_WP=

MOCK_VDAT_FLAGS=
MOCK_TPM_FWVER=
MOCK_FIRMWARE_INFO="
Key block:
  Size:                2232
  Flags:               7 (ignored)
  Data key algorithm:  7 RSA4096 SHA256
  Data key version:    2
  Data key sha1sum:    e2c1c92d7d7aa7dfed5e8375edd30b7ae52b7450
Preamble:
  Size:                  2164
  Header version:        2.1
  Firmware version:      1
  Kernel key algorithm:  7 RSA4096 SHA256
  Kernel key version:    1
  Kernel key sha1sum:    5d2b220899c4403d564092ada3f12d3cc4483223
  Firmware body size:    456411
  Preamble flags:        1
Body verification succeeded."

SAMPLE_FILE="../LICENSE"
SAMPLE_FILE_SIZE="1566"
SAMPLE_FILE_HASH="562c740877935f40b262db8af30bca36"

#----------------------------------------------------------------------
# Mock version of crossystem.
crossystem() {
  case "$1" in
    unittest)
      echo "ok"
      ;;
    unittest=set)
      echo "set ok"
      ;;
    query_invalid)
      echo "failure" >&2
      return ${FLAGS_FALSE}
      ;;
    fwupdate_tries)
      echo "1"
      ;;
    fwupdate_tries=2)
      echo "tries=2"
      ;;
    fwb_tries)
      echo 1
      ;;
    fwb_tries=2)
      echo "tries=2"
      ;;
    tpm_fwver)
      echo "${MOCK_TPM_FWVER}"
      ;;
    wpsw_cur)
      echo "${MOCK_WPSW_CUR}"
      ;;
    wpsw_boot)
      echo "${MOCK_WPSW_BOOT}"
      ;;
    vdat_flags)
      echo "${MOCK_VDAT_FLAGS}"
      ;;
  esac
}

# Mock version of flashrom
flashrom() {
  local opts="$*"
  if [ "${opts%--wp-status}" = "${opts}" ]; then
    die "Sorry, mock does not support $*." >&2
  fi
  local enabled="write protect is enabled"
  local disabled="write protect is disabled"
  if [ "${opts%-p host*}" != "${opts}" ]; then
    [ -n "${MOCK_MAIN_WP}" ] || enabled="${disabled}"
  elif [ "${opts%-p ec:dev=1}" != "${opts}" ]; then
    [ -n "${MOCK_PD_WP}" ] || enabled="${disabled}"
  elif [ "${opts%-p ec*}" != "${opts}" ]; then
    [ -n "${MOCK_EC_WP}" ] || enabled="${disabled}"
  else
    die "Sorry, invalid syntax for mock: $*"
  fi
  echo "$enabled"
}

#----------------------------------------------------------------------
test_cros_get_file() {
  assertEquals "$(cros_get_file_hash "${SAMPLE_FILE}")" "${SAMPLE_FILE_HASH}"
  assertEquals "$(cros_get_file_size "${SAMPLE_FILE}")" "${SAMPLE_FILE_SIZE}"
}

test_cros_compare() {
  assertEquals "$(cros_compare_version "Link.1.2.3" "Link.1.2.3")" "0"
  assertEquals "$(cros_compare_version "Link.12.2.3" "Link.1.2.3")" "1"
  assertEquals "$(cros_compare_version "Link.1.2.3" "Link.13.2.3")" "-1"

  assertTrue 'cros_version_greater_than "Link.1.2.3" "Link.1.1.1"'
  assertFalse 'cros_version_greater_than "Link.1.2.3" "Link.1.10.1"'

  assertTrue 'cros_compare_file ${SAMPLE_FILE} ${SAMPLE_FILE}'
  assertFalse 'cros_compare_file ${SAMPLE_FILE} $0'
}

test_cros_prop() {
  assertEquals "$(cros_get_prop unittest)" "ok"
  assertTrue "cros_set_prop unittest=set"
  assertEquals "$(cros_query_prop unittest)" "ok"
  assertEquals "$(cros_query_prop query_invalid)" ""
  assertTrue "cros_query_prop query_invalid"
}

test_cros_startup_update_tries() {
  assertEquals "$(cros_get_startup_update_tries)" "1"
  assertEquals "$(cros_set_startup_update_tries 2)" "tries=2"
}

test_cros_fwb() {
  assertEquals "$(cros_get_fwb_tries)" "1"
  assertEquals "$(cros_set_fwb_tries 2)" "tries=2"
}

test_write_protected() {
  MOCK_WPSW_CUR=
  MOCK_WPSW_BOOT=
  assertTrue "N,N" cros_is_hardware_write_protected
  MOCK_WPSW_BOOT=0
  assertFalse "N,0" cros_is_hardware_write_protected
  MOCK_WPSW_BOOT=1
  assertTrue "N,1" cros_is_hardware_write_protected
  MOCK_WPSW_CUR=1
  MOCK_WPSW_BOOT=0
  assertTrue "1,0" cros_is_hardware_write_protected
  MOCK_WPSW_CUR=0
  MOCK_WPSW_BOOT=1
  assertFalse "0,1" cros_is_hardware_write_protected

  MOCK_MAIN_WP=Y
  MOCK_EC_WP=Y
  MOCK_PD_WP=Y
  assertTrue "M" "cros_is_software_write_protected '${TARGET_OPT_MAIN}'"
  assertTrue "E" "cros_is_software_write_protected '${TARGET_OPT_EC}'"
  assertTrue "P" "cros_is_software_write_protected '${TARGET_OPT_PD}'"
  MOCK_MAIN_WP=
  MOCK_EC_WP=
  MOCK_PD_WP=
  assertFalse "!M" "cros_is_software_write_protected '${TARGET_OPT_MAIN}'"
  assertFalse "!E" "cros_is_software_write_protected '${TARGET_OPT_EC}'"
  assertFalse "!P" "cros_is_software_write_protected '${TARGET_OPT_PD}'"
}

test_cros_firmware_info() {
  # TODO(hungte) Test following functions:
  # cros_get_rw_firmware_info
  # cros_check_same_root_keys

  assertEquals "1" \
    "$(cros_get_firmware_preamble_flags "${MOCK_FIRMWARE_INFO}")"

  FLAGS_debug=$FLAGS_TRUE
  MOCK_TPM_FWVER=0x030003
  assertFalse 'cros_check_tpm_key_version "${MOCK_FIRMWARE_INFO}"'
  MOCK_TPM_FWVER=0x010002
  assertTrue 'cros_check_tpm_key_version "${MOCK_FIRMWARE_INFO}"'
  MOCK_TPM_FWVER=0x020001
  assertFalse 'cros_check_tpm_key_version "${$MOCK_FIRMWARE_INFO}"'

  MOCK_VDAT_FLAGS=$((0x1302))
  assertFalse cros_is_ro_normal_boot
  MOCK_VDAT_FLAGS=$((0x130f))
  assertTrue cros_is_ro_normal_boot
}

test_cros_check_stable_firmware() {
  FLAGS_update_main=${FLAGS_TRUE}
  FLAGS_update_ec=${FLAGS_FALSE}
  FLAGS_update_pd=${FLAGS_FALSE}
  RO_FWID="Link.123.0"
  STABLE_FWID="Link.123.1"
  FLAGS_wp=true
  assertTrue "wp=1,s" cros_check_stable_firmware
  FLAGS_wp=false
  assertFalse "wp=0,s" cros_check_stable_firmware
  RO_FWID="Link.123.2"
  assertTrue "wp=0,s2" cros_check_stable_firmware

  RO_FWID="Link.123.0"
  STABLE_FWID=""
  FLAGS_wp=true
  assertTrue "wp=1,n" cros_check_stable_firmware
  FLAGS_wp=false
  assertFalse "wp=0,n" cros_check_stable_firmware

  FLAGS_update_main=${FLAGS_FALSE}
  FLAGS_update_ec=${FLAGS_TRUE}
  FLAGS_update_pd=${FLAGS_FALSE}
  ECID="link/v1.123.0"
  STABLE_ECID="link/v1.123.1"
  FLAGS_wp=true
  assertTrue "wp=1,es" cros_check_stable_firmware
  FLAGS_wp=false
  assertFalse "wp=0,es" cros_check_stable_firmware
  ECID="link/v1.123.2"
  assertTrue "wp=0,es2" cros_check_stable_firmware
  ECID="link/v1.123.0"
  STABLE_ECID=""

  FLAGS_update_main=${FLAGS_FALSE}
  FLAGS_update_ec=${FLAGS_FALSE}
  FLAGS_update_pd=${FLAGS_TRUE}
  PDID="link/v1.123.0"
  STABLE_PDID="link/v1.123.1"
  FLAGS_wp=true
  assertTrue "wp=1,es" cros_check_stable_firmware
  FLAGS_wp=false
  assertFalse "wp=0,es" cros_check_stable_firmware
  PDID="link/v1.123.2"
  assertTrue "wp=0,es2" cros_check_stable_firmware
  PDID="link/v1.123.0"
  STABLE_PDID=""

  FLAGS_wp=true
  assertTrue "wp=1,en" cros_check_stable_firmware
  FLAGS_wp=false
  assertFalse "wp=0,en" cros_check_stable_firmware
}

test_cros_override_rw_firmware_by_version() {
  local min_ro="Google_Snow_2695.90.0"
  local need_update="Google_Snow_2695.87"
  local no_update="Google_Snow_2695.132.0"

  IMAGE_MAIN_RW=""
  FLAGS_mode=factory_install
  assertTrue "Modes without RW-only update" \
    "cros_override_rw_firmware_by_version ${need_update} ${min_ro}"

  FLAGS_mode=recovery
  FLAGS_wp=false
  assertTrue "Recovery without write protection should simply return." \
    "cros_override_rw_firmware_by_version ${need_update} ${min_ro}"
  FLAGS_wp=true
  FLAGS_mode=recovery
  assertTrue "Recovery+WP should pass if current >= min_ro." \
    "cros_override_rw_firmware_by_version ${no_update} ${min_ro}"
  assertFalse "Recovery+WP should fail if current < min_ro and no RW image" \
    "cros_override_rw_firmware_by_version ${need_update} ${min_ro}"

  FLAGS_wp=true
  FLAGS_mode=autoupdate
  assertTrue "Autoupdate+WP should pass if current >= min_ro." \
    "cros_override_rw_firmware_by_version ${no_update} ${min_ro}"
  assertFalse "Autoupdate+WP should fail if current < min_ro and no RW image" \
    "cros_override_rw_firmware_by_version ${need_update} ${min_ro}"
  FLAGS_wp=false
  assertTrue "Autoupdate should pass if current >= min_ro." \
    "cros_override_rw_firmware_by_version ${no_update} ${min_ro}"
  assertFalse "Autoupdate should fail if current < min_ro and no RW image" \
    "cros_override_rw_firmware_by_version ${need_update} ${min_ro}"
}

test_cros_check_compatible_platform() {
  assertTrue "platform check should pass if image = platform" \
    "cros_check_compatible_platform Google_Link Google_Link"
  assertFalse "platform check should fail if image != platform" \
    "cros_check_compatible_platform Google_Link Google_Samus"

  assertTrue "platform check should pass if image is empty" \
    "cros_check_compatible_platform '' Google_Samus"
  assertFalse "platform check should fail if target is empty" \
    "cros_check_compatible_platform Google_Link ''"

  assertTrue "platform check should pass on Snow image." \
    "cros_check_compatible_platform Google_Snow Google_Snow"
  assertTrue "platform check should pass on Snow image RevN boards." \
    "cros_check_compatible_platform Google_Snow Google_Snow_Rev4"
  assertFalse "platform check should fail on unknown Snow." \
    "cros_check_compatible_platform Google_Snow Google_Snow_New"
}

. ../lib/shunit2/shunit2
